# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import time
from pydantic import TypeAdapter
from tortoise.queryset import Q
from tortoise.transactions import in_transaction
from hypertext import PagingResult
from exception import AppException
from common.utils.urls import UrlUtil
from common.utils.tools import ToolsUtil
from common.utils.config import ConfigUtil
from common.enums.voice import VoiceEnum
from common.enums.agent import AgentEnum
from common.enums.models import ModelsEnum
from common.models.users import UserModel
from common.models.dev import DevModelsModel
from common.models.agents import AgentMcpModel
from common.models.agents import AgentKnowModel
from common.models.agents import AgentRobotModel
from common.models.agents import AgentSquareModel
from common.models.agents import AgentRobotTeamedModel
from common.models.agents import AgentRobotCollectModel
from common.service.clean_logic import CleanLogic
from common.service.subscribe_logic import SubscribeLogic
from plugins.censor.driver import CensorDriver
from apps.api.schemas.agent.mcp.client_schema import McpWholeVo
from apps.api.schemas.agent.know.manage_schema import KnowWholeVo
from apps.api.schemas.agent.robot import smarter_schema as schema
from apps.api.service.agent.robot.teamed_service import RobotTeamedService


class RobotSmarterService:
    """ 智能体服务类 """

    @classmethod
    async def lists(cls, params: schema.RobotSmarterSearchIn, user_id: int) -> PagingResult[schema.RobotSmarterListVo]:
        """
        智能体列表。

        Args:
            params (schema.RobotSmarterSearchIn): 智能体查询参数。
            user_id (int): 用户ID。

        Returns:
            PagingResult[schema.RobotSmarterListVo]: 智能体分页列表Vo。

        Author:
            zero
        """
        # 构建查询条件
        where = [Q(is_delete=0)]
        if params.keyword:
            where.append(Q(name__icontains=params.keyword))

        # 成员身份过滤
        permissions = {user_id: AgentEnum.POWER_OWNED}
        if params.owner == "private":
            where.append(Q(user_id=user_id))
        elif params.owner == "cooperate":
            team_lis = await AgentRobotTeamedModel.filter(user_id=user_id).all()
            team_ids = [item.robot_id for item in team_lis if item.robot_id]
            permissions = {team.user_id: team.power for team in team_lis}
            where.append(Q(id__in=team_ids or [0]))
        else:
            team_lis = await AgentRobotTeamedModel.filter(user_id=user_id).all()
            team_ids = [item.robot_id for item in team_lis if item.robot_id]
            if team_ids:
                permissions = {team.user_id: team.power for team in team_lis}
                where.append(Q(user_id=user_id) | Q(id__in=team_ids or [0]))
            else:
                where.append(Q(user_id=user_id))

        # 查询智能体列表
        _model = AgentRobotModel.filter(*where).order_by("-sort", "-update_time")
        _pager = await AgentRobotModel.paginate(
            model=_model,
            page_no=params.page_no,
            page_size=params.page_size,
            fields=["id", "code", "user_id", "logo", "name", "intro", "llm_models", "create_time", "update_time"],
            datetime_format="%m/%d %H:%M"
        )

        # 查询用户信息
        user_ids = [item["user_id"] for item in _pager.lists if item["user_id"]]
        users = await UserModel.queryUserDictsById(user_ids, ["id", "sn", "avatar", "nickname"])

        # 查询模型信息
        models = [item["llm_models"] for item in _pager.lists if item["llm_models"]]
        llm_models = await DevModelsModel.get_model_alias(models, ModelsEnum.TYPE_CHAT)

        # 查询广场分享
        robot_ids = [item["id"] for item in _pager.lists if item["id"]]
        squares = await AgentSquareModel.querySquareDictsByRobotId(robot_ids)

        # 循环处理列表
        lists = []
        for item in _pager.lists:
            item["logo"] = await UrlUtil.to_absolute_url(item["logo"])
            item["power"] = permissions.get(item["user_id"], 0)
            item["is_owner"] = True if item["user_id"] == user_id else False

            # 用户
            user = users.get(item['user_id'], {})
            item["user"] = {
                "sn": user.get("sn", ""),
                "nickname": user.get("nickname", ""),
                "avatar": await UrlUtil.to_absolute_url(user.get('avatar', ""))
            }

            # 模型
            item["models"] = llm_models.get(item["llm_models"], {})

            # 广场
            square = squares.get(item["id"], {})
            item["is_square"] = square.get("audit_status", -1)
            if item["is_square"] == 0 and square.get("is_enable") == 0:
                item["is_square"] = - 2

            # 序列化
            vo = TypeAdapter(schema.RobotSmarterListVo).validate_python(item)
            lists.append(vo)

        _pager.lists = lists
        return _pager

    @classmethod
    async def detail(cls, code: str, user_id: int) -> schema.RobotSmarterDetailVo:
        """
        智能体详情。

        Args:
            code (str): 智能体编号。
            user_id (int): 用户ID。

        Returns:
            schema.RobotSmarterDetailVo: 智能体详情Vo。

        Author:
            zero
        """
        # 查找机器人
        robot = await AgentRobotModel.filter(code=code, is_delete=0).get()

        # 权限的验证
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id, AgentEnum.POWER_VIEW)

        # 关联知识库
        knows = []
        if robot.know_ids:
            models = []
            know_ids = robot.know_ids.split(",")
            _knows = await AgentKnowModel.filter(id__in=know_ids, is_delete=0).all()
            vector = [item.vector_model.split(":")[0] for item in _knows if item.vector_model]
            if vector:
                models = await (DevModelsModel
                                .filter(scene=ModelsEnum.TYPE_VECTOR)
                                .filter(mark__in=list(set(vector)))
                                .filter(is_delete=0)
                                .all().values_list("mark", "name"))
                models = {k: v for k, v in models}

            for item in _knows:
                mark = item.vector_model.split(":")[0]
                knows.append(KnowWholeVo(
                    code=item.code,
                    logo=await UrlUtil.to_absolute_url(item.logo),
                    name=item.name,
                    vector_name=models.get(mark) or "-",
                    vector_model=item.vector_model,
                    is_disable=item.is_disable,
                ))

        # 关联MCP
        mcps = []
        if robot.mcp_ids:
            mcp_ids = robot.mcp_ids.split(",")
            _mcps = await AgentMcpModel.filter(id__in=mcp_ids, is_delete=0).all()
            for item in _mcps:
                mcps.append(McpWholeVo(
                    code=item.code,
                    logo=await UrlUtil.to_absolute_url(item.logo),
                    name=item.name,
                    is_disable=item.is_disable,
                ))

        # 模型的配置
        llm_name = ""
        llm_config = []
        if robot.llm_models:
            configs = json.loads(robot.llm_config)
            mark: str = robot.llm_models.split(":")[0]
            models = await DevModelsModel.filter(mark=mark, scene=ModelsEnum.TYPE_CHAT).first()
            if models:
                llm_name = models.name
                setting: dict = ModelsEnum.chat_models().get(models.channel)
                for conf in setting.get("configs")[0]["config"]:
                    conf["default"] = configs.get(conf["key"])
                    llm_config.append(conf)

        # 权限查询
        power: int = RobotTeamedService.POWER_OWNED
        if robot.user_id != user_id:
            power = await RobotTeamedService.has_permission(robot.code, user_id, RobotTeamedService.POWER_VIEW)

        # 语音播报
        broad = json.loads(robot.broadcast or "{}")
        voice = broad.get("model", "").split(":")[0]
        speech = VoiceEnum.get_default_speech(voice) or {}
        broadcast = {
            "model": broad.get("model", ""),
            "speed": broad.get("speed", 0.25),
            "speed_range": speech.get("speed_range", [0.25, 4.0, 0.1])
        }

        # 联网搜索
        is_web_search: int = int(await ConfigUtil.get("web_search", "status", 0))

        # 返回信息
        return schema.RobotSmarterDetailVo(
            code=robot.code,
            type=robot.type,
            logo=await UrlUtil.to_absolute_url(robot.logo),
            name=robot.name,
            intro=robot.intro,
            power=power,
            owned=True if robot.user_id == user_id else False,
            search_mode=robot.search_mode,
            search_tokens=robot.search_tokens,
            search_similar=robot.search_similar,
            ranking_status=robot.ranking_status,
            optimize_ask=robot.optimize_ask,
            annex_similar=robot.annex_similar,
            ranking_score=robot.ranking_score,
            ranking_model=robot.ranking_model,
            optimize_model=robot.optimize_model,
            prompt_roles=robot.prompt_roles or "",
            prompt_limit=robot.prompt_limit or "",
            welcome_text=robot.welcome_text or "",
            question_guide=robot.question_guide,
            max_upload_num=robot.max_upload_num,
            is_upload_files=robot.is_upload_files,
            is_upload_image=robot.is_upload_image,
            is_web_search=is_web_search,
            is_show_timed=robot.is_show_timed,
            is_show_quote=robot.is_show_quote,
            is_show_context=robot.is_show_context,
            is_original_reply=robot.is_original_reply,
            empty_search_text=robot.empty_search_text or "",
            empty_search_type=robot.empty_search_type or "",
            llm_name=llm_name,
            llm_models=robot.llm_models,
            llm_config=llm_config,
            knows=knows,
            mcps=mcps,
            broadcast=broadcast
        )

    @classmethod
    async def add(cls, post: schema.RobotSmarterAddIn, user_id: int):
        """
        智能体创建。

        Args:
            post (schema.RobotSmarterAddIn): 智能体新增参数。
            user_id (int): 用户ID。

        Author:
            zero
        """
        try:
            # 验证剩余数
            surplus: int = await SubscribeLogic.get_surplus_quota(user_id, scene="robot")
            if surplus <= 0:
                raise AppException("智能体创建额度不足!")

            # 查询大模型
            _models = await (DevModelsModel
                             .filter(scene=ModelsEnum.TYPE_CHAT)
                             .filter(is_disable=0, is_delete=0)
                             .order_by("sort", "-id")
                             .first())

            # 处理大模型
            llm_model = ""
            llm_config = {}
            if _models:
                config = json.loads(_models.config)
                models = json.loads(_models.models)
                config.pop("check_key", None)
                config.pop("agency_api", None)
                config.pop("global_directives", None)
                llm_config = config
                for llm in models:
                    if llm["status"]:
                        llm_model = f"{_models.mark}:{llm['model']}"
                        break

            # 违规审查 (信息)
            review = CensorDriver.REVIEW_TEXTS
            values = post.name + "\n" + post.intro
            censor_info = await CensorDriver.text(user_id=user_id, text=values, scene=review)
            if censor_info is not None:
                raise AppException(f"智能体信息违规: {censor_info}")

            # 创建数据
            code: str = await ToolsUtil.make_rand_sn(AgentRobotModel, 'code', 16, 'char')
            await AgentRobotModel.create(
                user_id=user_id,
                code=code.lower(),
                type="simple",
                logo="static/images/default_robot.png",
                name=post.name,
                intro=post.intro,
                llm_models=llm_model,
                llm_config=json.dumps(llm_config),
                search_mode="similar",
                search_tokens=5000,
                search_similar=0.4,
                ranking_status=0,
                optimize_ask=0,
                annex_similar=0.5,
                ranking_score=0.4,
                ranking_model="",
                prompt_roles="",
                prompt_limit="",
                welcome_text="",
                question_guide=0,
                empty_search_text="",
                empty_search_type="ai",
                max_upload_num=0,
                is_upload_files=0,
                is_upload_image=0,
                is_show_timed=1,
                is_show_quote=1,
                is_show_context=1,
                is_original_reply=0,
                create_time=int(time.time()),
                update_time=int(time.time())
            )
        except Exception as e:
            raise AppException(str(e))

    @classmethod
    async def edit(cls, post: schema.RobotSmarterEditIn, user_id: int):
        """
        智能体编辑。

        Args:
            post (schema.RobotSmarterEditIn): 智能体编辑参数。
            user_id (int): 用户ID。

        Author:
            zero
        """
        # 智能体校验
        robot = await AgentRobotModel.filter(code=post.code, is_delete=0).get()

        # 权限的校验
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id)

        # 大模型校验
        mark = post.llm_models.split(":")[0]
        models = await DevModelsModel.filter(mark=mark, is_delete=0).first()
        if not models or models.is_disable:
            raise AppException("模型已停用")

        # 关联知识库
        know_ids = []
        if post.knows:
            knows = await AgentKnowModel.filter(code__in=post.knows, is_delete=0).all().values("id")
            know_ids = [str(item["id"]) for item in knows if item["id"]]
            if len(know_ids) != len(post.knows):
                raise AppException("关联知识库存在丢失,请刷新页面重新选择")

        # 关联MCP
        mcp_ids = []
        if post.mcps:
            mcps = await AgentMcpModel.filter(code__in=post.mcps, is_delete=0).all().values("id")
            mcp_ids = [str(item["id"]) for item in mcps if item["id"]]
            if len(mcp_ids) != len(post.mcps):
                raise AppException("关联MCP存在丢失,请刷新页面重新选择")

        # 违规审查
        review = CensorDriver.REVIEW_TEXTS
        values = post.name + "\n" + post.intro
        censor_info = await CensorDriver.text(user_id=user_id, text=values, scene=review)
        if censor_info is not None:
            raise AppException(f"智能体信息违规: {str(censor_info)}")

        # 更新智能体
        await AgentRobotModel.filter(id=robot.id).update(
            user_id=user_id,
            know_ids=",".join(know_ids),
            mcp_ids=",".join(mcp_ids),
            logo=UrlUtil.to_relative_url(post.logo),
            name=post.name,
            intro=post.intro,
            llm_models=post.llm_models,
            llm_config=json.dumps(post.llm_config),
            search_mode=post.search_mode,
            search_tokens=post.search_tokens,
            search_similar=post.search_similar,
            ranking_status=post.ranking_status,
            optimize_ask=post.optimize_ask,
            annex_similar=post.annex_similar,
            ranking_score=post.ranking_score,
            ranking_model=post.ranking_model,
            optimize_model=post.optimize_model,
            prompt_roles=post.prompt_roles,
            prompt_limit=post.prompt_limit,
            welcome_text=post.welcome_text,
            question_guide=post.question_guide,
            empty_search_text=post.empty_search_text,
            empty_search_type=post.empty_search_type,
            max_upload_num=post.max_upload_num,
            is_upload_files=post.is_upload_files,
            is_upload_image=post.is_upload_image,
            is_show_timed=post.is_show_timed,
            is_show_quote=post.is_show_quote,
            is_show_context=post.is_show_context,
            is_original_reply=post.is_original_reply,
            update_time=int(time.time()),
            broadcast=json.dumps({
                "model": post.broadcast.get("model", ""),
                "speed": post.broadcast.get("speed", 1)
            }, ensure_ascii=False)
        )

    @classmethod
    async def delete(cls, code: str, user_id: int):
        """
        智能体删除。

        Args:
            code (str): 智能体编号。
            user_id (int): 用户ID。

        Author:
            zero
        """
        # 智能体验证
        robot = await AgentRobotModel.filter(code=code, is_delete=0).get()

        # 权限的校验
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id, RobotTeamedService.POWER_OWNED)

        async with in_transaction("mysql"):
            # 智能体删除
            await AgentRobotModel.filter(code=code).update(
                is_delete=1,
                delete_time=int(time.time())
            )
            # 关联表清理
            await CleanLogic.clean_robot(robot.id)

    @classmethod
    async def clone(cls, code: str, user_id: int):
        """
        智能体克隆。

        Args:
            code (str): 智能体编号。
            user_id (int): 用户ID。

        Author:
            zero
        """
        # 智能体验证
        robot = await AgentRobotModel.filter(code=code, is_delete=0).get()

        # 权限的校验
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id, RobotTeamedService.POWER_OWNED)

        # 验证剩余数
        surplus: int = await SubscribeLogic.get_surplus_quota(user_id, scene="robot")
        if surplus <= 0:
            raise AppException("智能体创建额度不足!")

        # 智能体克隆
        code: str = await ToolsUtil.make_rand_sn(AgentRobotModel, 'code', 16, 'char')
        robot.code = code
        robot.name = robot.name + " - 副本"
        robot.create_time = int(time.time())
        robot.update_time = int(time.time())
        robot = robot.__dict__
        del robot["id"]
        await AgentRobotModel.create(**robot)

    @classmethod
    async def square(cls, code: str, user_id: int):
        """
        智能体公开。

        Args:
            code (str): 智能体编号。
            user_id (int): 用户ID。

        Author:
            zero
        """
        # 智能体验证
        robot = await AgentRobotModel.filter(code=code, is_delete=0).get()

        # 权限的校验
        if robot.user_id != user_id:
            await RobotTeamedService.has_permission(robot.code, user_id, RobotTeamedService.POWER_ADMIN)

        # 加入/退出广场
        square = await AgentSquareModel.filter(robot_id=robot.id).first()
        if square is None:
            await AgentSquareModel.create(
                cate_id=0,
                robot_id=robot.id,
                audit_status=1,
                is_enable=1,
                create_time=int(time.time()),
                update_time=int(time.time())
            )
        else:
            if square.is_delete == 1:
                await AgentSquareModel.filter(id=square.id).update(
                    is_delete=0,
                    delete_time=0
                )
            else:
                await AgentSquareModel.filter(id=square.id).update(
                    is_delete=1,
                    delete_time=int(time.time())
                )

    @classmethod
    async def collect(cls, code: str, user_id: int):
        # 验证智能体
        robot = await AgentRobotModel.filter(code=code).first()
        if robot is None:
            raise AppException("智能体不存在!")
        if robot.is_delete:
            raise AppException("智能体已被删!")

        # 查询收藏
        collect = await (AgentRobotCollectModel
                         .filter(user_id=user_id)
                         .filter(robot_id=robot.id)
                         .first())

        # 查询广场
        # square = await AgentSquareModel.filter(robot_id=robot.id, is_delete=0).first()
        # if square is None and (collect is None or robot.is_delete):
        #     raise AppException("当前智能体尚未公开不可收藏")

        # 处理收藏
        if collect is None:
            await AgentRobotCollectModel.create(
                user_id=user_id,
                robot_id=robot.id,
                total_num=1,
                create_time=int(time.time()),
                update_time=int(time.time())
            )
        else:
            if collect.is_delete:
                await AgentRobotCollectModel.filter(user_id=user_id).update(
                    total_num=collect.total_num + 1,
                    is_delete=0,
                    delete_time=0
                )
            else:
                await AgentRobotCollectModel.filter(user_id=user_id).update(
                    is_delete=1,
                    delete_time=int(time.time())
                )
